昨天把 torch.compile 的 CPU pipeline 走過一次，Python 程式先由 TorchDynamo 擷取成一張或多張 FX graph，TorchInductor 再替可編譯的區段產生 C++，最後交給目標機器上的 compiler 做成 shared object。

前幾天走的是手寫 SGLang RVV backend，挑出重要的運算子，自己寫 RVV kernel、dispatch 與 Python binding。這條路很適合先把特定 hot path 加速，不過每多支援一種模型結構或運算子，就要繼續維護 runtime 端，那能不能減少手寫 RVV 運算子? 部分交給 TorchInductor 來完成?

原本 TorchInductor 還沒支援 RVV，把 RVV 接進 TorchInductor 好處是可以把這些硬體相關工作放進 PyTorch compiler backend。上層模型維持一般的 F.linear 寫法，Inductor 根據 shape、dtype、layout 與 CPU capability 選擇 kernel。不同模型只要進入相同的 Linear path，就能共用同一套 codegen 實作和 fallback policy，不過實際產生的 compiled artifact 可能因為 shape、guards 和 weight 條件而不同。

接下來二天，我會用最近剛投 PyTorch Conference North America 2026 的 poster，題目是 TorchInductor for SGLang Inference on RISC-V CPUs 當主線，拆解這個專案的動機、實作與效能證據。
我的實驗先處理 LLM 裡反覆執行的 Linear，再用 SGLang workload 檢查 prefill、decode、編譯時間和記憶體代價。端到端收益取決於支援的 shape、weight layout 與實際 serving 流程，後面的實驗會把這些條件一起列出來。

本篇大綱

• 先說明 poster 想處理的 PyTorch 支援缺口
• 說清楚這次 RVV codegen 位在 compiler pipeline 的哪一段
• 再看 full CPU graph capture 的內容、成本，以及為什麼改用 Linear regional boundary
• 用 P1–P3 拆開 ISA selection、microkernel 與 template routing
• 讀 P1 的 VecRVV、compiler flags 和 capability selection
• 讀 P2 的 M=1 GEMV 與 M>=2 GEMM
• 讀 P3 的靜態 shape gate 和 fallback
• 最後整理這三段如何從 F.linear 走到 RVV C++

poster 想處理的 PyTorch 支援缺口

紅色框的位置就是這次工作的範圍，初期讓一般 PyTorch Linear 經過 torch.compile 後，能在支援的條件下產生 RVV BF16 GEMV / GEMM。

SGLang 是當作提供 LLM serving workload、model boundary 與 benchmark entry point，實作重點放在 TorchInductor 的 RVV code generation。

這次的 codegen 做到哪一層

這裡的 RVV codegen，目前主要做在 TorchInductor 的 C++ template 與 microkernel 層。Inductor 先把符合條件的 Linear route 到 CppGemmTemplate，我加入的 RVV microkernel 再產生包含 riscv_vector.h intrinsic 的 C++，例如 __riscv_vle16_v_u16m2 和 __riscv_vfmacc_vf_f32m4。最後由 GCC 把這些 intrinsic 編成 RVV 指令。

目前實作沿用 TorchInductor 已有的 graph、lowering 和 template 架構，新增的是 RVV target selection、Linear routing 與 target-specific kernel codegen。這個階段還沒有新增一套 IR，也沒有實作跨多層 IR 的 rewrite、算子融合或全圖 scheduling 最佳化。目前是讓既有 Inductor pipeline 能替支援的 BF16 Linear 產生並選用 RVV kernel；更廣的 IR-level optimization 可以作為後續方向。

這個缺口會直接出現在 LLM inference。Transformer block 裡的 q_proj、k_proj、v_proj，或某些實作合併後的 qkv_proj，加上 o_proj、gate_up_proj、down_proj 和最後的 lm_head，都會執行大型 Linear。

P2 會再把 Linear 拆成 M、N、K，用 decode 與 prefill 的實際 shape 解釋為什麼需要兩種 microkernel。

Full CPU graph 帶來的下一個問題

這裡的 full CPU graph 是 SGLang 對 compile boundary 的命名：它把整個 model.forward 交給 torch.compile，再替指定的 batch size 保存可重複呼叫的 callable。目前 CPU graph runner 使用 torch.compile(..., dynamic=False)，沒有設定 fullgraph=True。

這兩個名字回答不同問題

• SGLang full CPU graph：torch.compile 從哪個函式入口開始？答案是整個 model.forward。
• torch.compile(fullgraph=True)：這個函式能不能被 Dynamo 擷取成單一 FX graph？如果中間出現 graph break，是否直接報錯？

目前使用預設的 fullgraph=False。如果 model.forward 內出現 graph break，Dynamo 可能把它切成多個 compiled region，中間回到 Python 執行。 所以 whole-model boundary 不等於一定只有一張 FX graph。

Full CPU graph capture 了哪些內容

CPUGraphRunner 會準備固定 shape 的輸入與 ForwardBatch，再呼叫包住 model.forward 的 compiled callable。Dynamo 擷取能編譯的 tensor operations，並替 shape、dtype、stride 和控制路徑建立 guards。dynamic=False 讓 batch size 或 token shape 的變化可能產生另一個 specialization；它不控制 graph break。

固定 shape 的 decode inputs
  -> whole-model model.forward boundary
  -> 一張或多張 Dynamo FX graphs + guards
  -> Inductor lowering 與 generated C++
  -> compiled .so 與可再次呼叫的 torch.compile wrapper

實驗也啟用 Inductor freezing。符合 capture 與 template 條件的穩定 Linear weight，可以轉成 microkernel 使用的 blocked layout。Generated .so 能透過 artifact cache 重用；packed tensor 通常仍是 process-local state。明天會再處理 packed weight 的 ownership。

為什麼 capture 成本高

CPUGraphRunner.capture() 的計時包含 Dynamo tracing 與 guards、Inductor lowering/codegen、符合條件的 freezing/prepack，以及 generated C++ 的 build、load 和第一次 warmup。Banana Pi 必須在板子上完成這些工作，因此 capture time 是整段 model-level compilation 的 process-wall cost；內部 timer 可能重疊，不能直接相加。

Banana Pi BPI-F3 上的 full CPU graph capture 成本很高。過去歷史實驗中，regional no-graph 的 process wall 是 592 秒；full CPU graph 是 1665 秒，其中 capture 約 820 秒。warm decode 從 1.320 tok/s 增加到 1.405 tok/s，約 6.4%。

在這次成功 capture 的路徑裡，Full CPU graph 可以重用 generated kernel 和 blocked weight，並減少 compiled regions 內的 Python / dispatcher 工作。不過一個 warmup input 無法涵蓋所有 batch size、shape 與控制路徑。完整模型 capture 不適合當目前 Banana Pi 實驗的主要 serving 啟動路徑，所以這裡改成只編譯計算量大的 Linear region。

所以選擇使用較窄的 regional torch.compile boundary

SGLang model
 -> 只將支援的 Linear region 交給 torch.compile
 -> TorchInductor 產生 RVV C++
 -> module 保留 explicit packed weight
 -> 其他 shape / dtype / layout 維持 fallback

regional boundary 避開 full-model capture，也要處理靜態 shape 造成的 artifact 數量，以及 runtime packed weight 的 ownership，這些留到明天。

今天先看 P1–P3

完整差異可以看 PyTorch compare page。今天依序讀三個固定 commit：P1 接上 RVV ISA selection、P2 加入 BF16 GEMV/GEMM、P3 將支援的靜態 shape route 進 template。P4–P6 留到 Day09。

PR #175746 之後，Inductor 還缺什麼

PyTorch PR #175746 提出，目前還沒有合併，主要想做是替 ATen CPU vector layer 加入 RVV Vectorized<T> 與 CPU capability 基礎，使用 Vectorized<T> 且有對應 RVV implementation 的 ATen operator，才有機會從這層支援受益，我的實作是基於這個 PR 基礎之上。

TorchInductor 有自己的 ISA 選擇與 C++ template 路徑，只套用 PR #175746 時，我在 baseline 看到

inductor_vec_isa = INVALID_VEC_ISA
inductor_cpp_march = None
generated pointwise kernel = scalar loop

這表示 ATen 已有 RVV foundation，但 Inductor 還沒有把 RVV 放進候選 ISA，generated C++ 也拿不到對應 flags、macros 與 Linear microkernel。P1 先補上這段連接。

P1：用 VecRVV 接上 ISA selection

P1 在 torch/_inductor/cpu_vec_isa.py 新增 VecRVV。完整修改可以看 P1 commit

class VecRVV(VecISA):
    _bit_width = 256
    _macro = [
        "CPU_CAPABILITY_RVV",
        "CPU_CAPABILITY=RVV",
        "HAVE_RVV_CPU_DEFINITION",
    ]
    _arch_flags = "-march=rv64gcv_zvl128b"
    _dtype_nelements = {
        torch.float: 8,
        torch.bfloat16: 16,
        torch.float16: 16,
    }

_bit_width = 256 是 Inductor 對這個實驗 target 使用的 compile-time vector width model。它要和實際機器能力一致，但只看這行程式碼，不能反過來證明硬體 VLEN。RVV 規格允許不同 VLEN；目前 prototype 先處理這條 256-bit 實驗線，multi-VLEN 仍是後續工作。

_arch_flags 會進到 generated C++ 的編譯命令。rv64gc 沒有 V extension；rv64gcv_zvl128b 才告訴 compiler 可以產生 RVV，且最小 VLEN 為 128 bit。

_macro 讓 generated C++ 和 PyTorch header 看到同一個 CPU capability。少了這一段，即使 compiler 接受 -march，include 進來的 vector implementation 仍可能走錯條件編譯分支。

_dtype_nelements 則告訴 Inductor：以目前 256-bit 設定估算，一個 vector 可容納 8 個 FP32 或 16 個 BF16 / FP16 element。

接著把 VecRVV 放進候選清單

elif arch == "riscv64":
    vec_rvv = VecRVV()
    if torch.backends.cpu.get_cpu_capability() == "RVV" and vec_rvv:
        isa_list.append(vec_rvv)

這裡同時檢查 machine architecture 與 ATen 回報的 CPU capability。

實驗環境會設定

export ATEN_CPU_CAPABILITY=rvv

這是實驗用的 capability override，不是硬體偵測工具。只有確認目標 CPU 和 toolchain 支援 RVV 時才能這樣設定；否則 generated binary 可能在執行時遇到 illegal instruction。

設定通過 architecture 與 capability 檢查後，valid_vec_isa_list() 會加入 VecRVV，Inductor 才能把 RVV macros 與 -march 帶進 generated C++ 的編譯流程。

P1 之後，pointwise codegen 已經能選到 RVV。不過 LLM 推論的大量時間在 Linear，下一步需要替 GEMV / GEMM template 加入 RVV microkernel。

P2：為什麼要拆成 M=1 和 M>=2

先把 Linear 寫成矩陣形狀

X [M, K] × Wᵀ [K, N] -> Y [M, N]

在 LLM 推論裡，M 會隨階段改變

M=1 時只有一列 activation，計算和資料重用方式接近 matrix-vector multiplication。M>=2 可以讓多列 activation 共用載入的 weight vector。若只準備一種 generic kernel，它可能無法同時適合這兩種 shape，所以 P2 分開註冊兩個版本。

P2 註冊兩個 class，完整修改可以看 P2 commit：

(1, 32, 1) 和 (4, 32, 1) 分別是 (block_m, block_n, block_k)。兩個 kernel 都沿 N 維一次處理 32 個輸出 channel；M>=2 版本一次保留最多 4 列 accumulator。

兩個 extra_check 分別接受 M=1 與 M>=2，而且都排除 dynamic M。目前 prototype 先處理可驗證的靜態形狀，dynamic shape 留在 fallback。

BF16 如何用 RVV 做 FP32 accumulate

目前 kernel 會把 BF16 bit pattern 展開成 FP32，再用 FP32 FMA 累加。M>=2 kernel 裡的 helper

auto load_bf16_as_f32 = [](const at::BFloat16* ptr, size_t vl) {
    const auto bf16_vec = __riscv_vle16_v_u16m2(
        reinterpret_cast<const uint16_t*>(ptr), vl);
    auto bits_vec = __riscv_vzext_vf2_u32m4(bf16_vec, vl);
    bits_vec = __riscv_vsll_vx_u32m4(bits_vec, 16, vl);
    return __riscv_vreinterpret_v_u32m4_f32m4(bits_vec);
};

BF16 保留 FP32 的高 16 bit。這段程式先載入 16-bit element，zero-extend 成 32 bit，再左移 16 bit，最後把相同 bit pattern reinterpret 成 FP32。

內層迴圈載入一段 weight，讓多列 activation 共用。下面只列出 acc0；其他 M row 以相同方式更新各自的 accumulator。

for (int64_t k = 0; k < K; ++k) {
    const auto weight_vec = load_bf16_as_f32(
        B + k * ldb + (n + j) * ldb_col, vl);

    acc0 = __riscv_vfmacc_vf_f32m4(
        acc0,
        static_cast<float>(A[(m + 0) * lda + k]),
        weight_vec,
        vl);
}

這裡的 vfmacc_vf 是 vector-scalar fused multiply-accumulate

acc[0:vl] += activation_scalar × weight_vector[0:vl]

weight 沿 N 維連續，activation 從不同 M row 取 scalar。同一個 weight_vec 可以更新 acc0、acc1、acc2、acc3，這是 M>=2 版本主要的資料重用。

M=1 kernel 也保留兩條路

• ldb_col == 1：weight 已是適合沿 N 讀取的 layout，沿 N 做 vector FMA。
• ldb_col != 1：row-major / transposed view 無法沿 N 連續讀取，改沿 K 做 dot product 和 reduction。

第二條是正確性與 runtime-weight fallback；明天的 explicit packed state 會讓常用路徑使用第一種 layout。

P3：只 route 有收益的靜態 BF16 shape

microkernel 註冊後，Inductor 還要決定何時選它。P3 在預設模式加入一組 shape gate。完整修改可以看P3 commit

_RVV_BF16_TEMPLATE_MIN_OPS = 1_000_000
_RVV_BF16_TEMPLATE_MAX_M = 128
_RVV_BF16_TEMPLATE_MIN_N = 1024
_RVV_BF16_TEMPLATE_MIN_K = 1024


def _is_profitable_rvv_bf16_template_shape(m, n, k, ops):
    return (
        ops >= _RVV_BF16_TEMPLATE_MIN_OPS
        and m <= _RVV_BF16_TEMPLATE_MAX_M
        and n >= _RVV_BF16_TEMPLATE_MIN_N
        and k >= _RVV_BF16_TEMPLATE_MIN_K
    )

這些門檻把路徑限制在 LLM projection-sized BF16 matmul

• N、K 至少 1024，排除大量小矩陣
• M 不超過 128，先聚焦 decode 和小段 prefill
• M × N × K 至少一百萬，排除較可能被 template 編譯與呼叫成本蓋過收益的小工作量
• M、N、K 必須能取得靜態 hint

當 shape 通過檢查，而且 create_micro_gemm() 選到 RVV BF16 class，tuned_mm() 會加入 C++ template choice。可以把 P3 commit 的 routing 簡化成：

if shape_passes_gate and rvv_microkernel_is_available:
    choices.append(CppGemmTemplate)
else:
    keep_existing_fallback()

這段讓 RVV static template 不必要求使用者先打開全域 max-autotune。沒有通過 gate 的 shape 仍保留 ATen external kernel 或其他既有 fallback；這不一定代表 graph break，也不等於整個函式回到 eager。P3 沒有把所有 CPU matmul 都強迫改走新 kernel。

P3 先限制「哪些 shape 值得編」。到了 SGLang serving，M=2、M=3、M=4 如果各自編一份 artifact，形狀數仍會增加。P5 會再用 2/4/8/16 bucket 把實際編譯版本壓到有限集合。

從 F.linear 走到 RVV C++

把 P1–P3 串起來，一個符合條件的 F.linear 會走這條路

F.linear(input, weight)
 -> aten.linear / aten.mm
 -> tuned_mm()
 -> use_cpp_gemm_template()
 -> static BF16 shape gate
 -> create_micro_gemm()
 -> M=1: CppMicroGemmRVVBF16M1
    M>=2: CppMicroGemmRVVBF16MGe2
 -> CppGemmTemplate 產生 C++
 -> VecRVV 提供 macros / -march
 -> GCC 15 編譯 shared object

實驗驗證使用一個乾淨獨立的 TORCHINDUCTOR_CACHE_DIR，再檢查 generated artifact

grep -R -l '__riscv_vfmacc' "${TORCHINDUCTOR_CACHE_DIR}" --include='*.cpp'
grep -R -n 'extern_kernels.mm' "${TORCHINDUCTOR_CACHE_DIR}" --include='*.py'
find "${TORCHINDUCTOR_CACHE_DIR}" -type f -name '*.so' -print

三個指令分別檢查 RVV intrinsic、ATen mm fallback 與 compiled .so。乾淨 cache 可以避免混入其他實驗，正式驗證還要 disassemble .so，不能只靠字串判定。

今天先走到這裡

• Full CPU graph 使用 whole-model model.forward boundary，但沒有設定 fullgraph=True；capture time 也包含 codegen、C++ build/load、符合條件的 prepack 與 warmup。
• 目前的 RVV codegen 位在 C++ template / microkernel 層，由 template 產生 RVV intrinsic C++；這一版尚未加入新的 RVV IR 或跨算子、跨 IR 的最佳化 pass。
• PR #175746 提供 ATen RVV foundation；P1 再把 RVV 接進 Inductor 的 ISA selection、macros 和 compiler flags。
• P2 用 M=1 GEMV 處理 decode，M>=2 GEMM 讓多列 activation 共用 weight vector，兩者都使用 FP32 accumulator。
• P3 只 route 靜態、夠大的 BF16 projection shape，其他情況保留 ATen fallback；下一篇再用 serving bucket 限制 artifact 數量。

明天會處理剩下 P4–P6。重點從 compiler template 移到 runtime ownership：packed weight 要由誰持有、weight reload 後怎麼避免讀到舊資料、兩條 decode sequence 形成 M=2 時怎麼 route，以及另一台 Banana Pi 要如何重現同一個環境和測試。

參考資料

• PyTorch PR #175746：RISC-V Vector Extension support for ATen
• PyTorch main...rvv_pytorch_inductor 完整比較
• P1 commit：enable RVV CPU vectorization
• P2 commit：add RVV BF16 GEMM and GEMV microkernels
• P3 commit：route static BF16 Linear shapes to RVV templates
• PyTorch rvv_pytorch_inductor branch
• SGLang issue #18072：2026 Q3–Q4 roadmap 與結果
• PyTorch tutorial：Introduction to torch.compile